/*
 * WebotsSupervisor.cpp
 *
 *  Created on: 29 Jun 2012
 *      Author: bweel
 */
#include <boost/lexical_cast.hpp>
#include <boost/functional/hash.hpp>
#include <boost/program_options.hpp>
#include <string>
#include <iostream>
#include <math.h>
#include <numeric>

#include "WebotsSupervisor.h"
#include "NNController.h"
#include "EggController.h"
#include "EvolutionaryController.h"
#include "AvoidanceController.h"
#include "Mathlib.h"

WebotsSupervisor::WebotsSupervisor (int argc, char **argv, RobotType *robot) : DemoSupervisor(argc, argv),
    robot(robot),
	startTime(0) {
}

bool WebotsSupervisor::initialise(int argc, char **argv) {
	startTime = getTime();
	worldModel.time = getTime();
	worldModel.robotId = getName();

	size_t modelIdx = worldModel.robotId.find_first_of(':');
	worldModel.robotModel = worldModel.robotId.substr(0,modelIdx);
	worldModel.robotNumber = boost::lexical_cast<int>(worldModel.robotId.substr(modelIdx+1,std::string::npos));
    
    if(!checkSupport() || !parseParameters(argc,argv)){
        return false;
    }
    
	initialiseRobotHardware();
    
	loadWorldModel();

	boost::property_tree::ptree supervisorTree;
	supervisorTree.put("state", 0);
	worldModel.additionalProperties.add_child("Supervisor", supervisorTree);

	worldModel.changePosAsked = false;
	worldModel.desiredHingePosition = 0;
	worldModel.desiredSpeed[0] = 0;
	worldModel.desiredSpeed[1] = 0;

	worldModel.organismChannel = -1;

    switch(controlType){
        case 0:
        default:
            control = new NNController(&worldModel, robot);
            if(loadNN){
                ((NNController*)control)->loadFromFile(NNFile);
            }
            break;
        case 1:
            control = new EggController(&worldModel, robot);
            break;
        case 2:
            control = new EvolutionaryController(&worldModel, robot,new NNController(&worldModel,robot));
            break;
        case 3:
            control = new AvoidanceController(&worldModel,robot);
            break;
    }
    
    return true;
}

bool WebotsSupervisor::checkSupport(){
    std::vector<std::string> supported;
    supported.push_back("KIT");
    supported.push_back("AW");
    supported.push_back("KIT_noWheel");
    supported.push_back("SimpleKIT");
    supported.push_back("Scout");
    supported.push_back("Scout_noWheel");
    
	if(std::find(supported.begin(),supported.end(),worldModel.robotModel) == supported.end()){
		std::cerr << "Error: robot model " <<  worldModel.robotModel << " is not supported! Robot model should be one of: ";
        for(int i=0;i<supported.size();i++){
            std::cerr << supported[i];
            if(i<supported.size()-1){
                std::cerr << ", ";
            }
        }
        std::cerr << ". Please use robot names as: <robot model>:<robot number>" << std::endl;
        return false;
	}
    return true;
}

bool WebotsSupervisor::parseParameters(int argc, char **argv){
    using namespace boost::program_options;
    options_description options("Allowed Options");
    options.add_options()
        ("help", "produce help message")
        ("egg","Whether the robot is an egg")
        ("controllerType", value<int>(&controlType)->required(),"The type of controller")
        ("loadNN", value<std::string>(&NNFile),"Load a Neural Network from a file")
        ("seed", value<int>(),"The random seed");
    
    variables_map vm;
    try {
        store(parse_command_line(argc,argv, options), vm);
        notify(vm);
    }catch(boost::program_options::required_option& e){
        std::cerr << worldModel.robotId << " " << e.what() << "\n";
        std::cerr << options << std::endl;
        return false;
    }
    
    if (vm.count("help")){
        std::cout << options << std::endl;
    }
    
    if(vm.count("egg")){
        worldModel.isEgg = true;
    }else{
        worldModel.isEgg = false;
    }
    
    if(vm.count("seed")) {
        srand(vm["seed"].as<int>());
        Mathlib::setRandSeed(vm["seed"].as<int>());
    }
    
    if(vm.count("loadNN")){
        loadNN = true;
    }
    
    return true;
}

void WebotsSupervisor::initialiseRobotHardware() {

//	std::cout << "Initializing robot harware...";
    robot->initialiseHardware();
}

WebotsSupervisor::~WebotsSupervisor() {
}

void WebotsSupervisor::controlStep() {
		loadWorldModel();

		boost::property_tree::ptree& supervisorTree = worldModel.additionalProperties.get_child("Supervisor");
		if(supervisorTree.get<int>("state") == 0) {
			if(getTime() - startTime < (2.0*(double)TIMESTEP/1000.0)){
				// wait for the first few steps
				return;
			}else{
				// Now we should have the correct presence stuff
				for (unsigned int i = 0; i < robot->getNumConnectors(); i++) {
					worldModel.connected[i] = detectPresence(i); // isConnected(i);
				}
                
                if(numConnected() == 0){
                    supervisorTree.put("state",3);
                }else{
                    supervisorTree.put("state",1);
                }
				startTime = getTime();
			}
		} else if (supervisorTree.get<int>("state") == 1) {
//            std::cout << worldModel.robotId << "Setting organism channel" << std::endl;
			setOrganismChannel();
		} else if(supervisorTree.get<int>("state") == 2) {
//            std::cout << worldModel.robotId << "Calculating organism size" << std::endl;
			calculateOrganismSize();
		} else if(supervisorTree.get<int>("state") == 3) {
//            std::cout << worldModel.robotId << "Controlling robot" << std::endl;
            control->step();
		}
		applyWorldModel();
}

void WebotsSupervisor::setLEDs(std::vector<LED*> toset, int colour) {
	for (int i = 0; i < toset.size(); i++) {
		toset[i]->set(colour);
	}
}

void WebotsSupervisor::setLEDColour(std::pair<int,int> ledColour){
    if(robot->hasLED(ledColour.first)){
//        std::cout << "Setting led " << ledColour.first << " to value " << ledColour.second << std::endl;
        robot->getLED(ledColour.first)->set(ledColour.second);
    }
}

void WebotsSupervisor::calculateOrganismSize(){
	static bool initialized = false;
	static bool sent = false;
	static bool done = false;

	if(!initialized) {
		for(int i=0;i<4;i++){
			if(!worldModel.connected[i]){
				counted[i] = 0;
			}else{
				counted[i] = -1;
			}
		}
		initialized = true;
	}

	// Check for ack
	for(size_t i=0;i<worldModel.msgReceived.size();i++){
		std::string message = worldModel.msgReceived[i].getDataString();

		if(message.find("ack") != std::string::npos){
			sent = true;
		}
	}

	if(numNotReceived(counted, 4) == 0){
		// we are root
		worldModel.organismSize = sum(counted,4) + 1;
		std::string msg = "organsimsize:"+boost::lexical_cast<std::string>(worldModel.organismSize);
		worldModel.msgToSend.push_back(new Message("+", msg));

		std::cout << worldModel.robotId << " found organism size: " << worldModel.organismSize << std::endl;
		boost::property_tree::ptree& supervisorTree = worldModel.additionalProperties.get_child("Supervisor");
		startTime = getTime();
		supervisorTree.put("state", 3);

		done = true;
	}else if(!sent && !done){
		if(numConnected() == 1){
			std::string address = getOpenConnection();
//			std::cout << worldModel.robotId << "I have one connection, sending message to " << address << std::endl;
			std::string msg = "treesize:1";
			worldModel.msgToSend.push_back(new Message(address,msg));
		}else{
			if (numNotReceived(counted, 4) == 1){
				// we are ready to send
				std::string address = getOpenConnection();

				int size = sum(counted,4) + 1;
//				std::cout << worldModel.robotId << "I have one connection, sending message to " << address << " current count: " << size << std::endl;

				std::string msg = "treesize:"+boost::lexical_cast<std::string>(size);
				worldModel.msgToSend.push_back(new Message(address, msg));
			}
			// wait till we have more
//			std::cout << worldModel.robotId << " waiting for more size messages" << std::endl;
		}
		for(size_t i=0;i<worldModel.msgReceived.size();i++){
			std::string message = worldModel.msgReceived[i].getDataString();

//			std::cout << worldModel.robotId << " received from: " << worldModel.msgReceived[i]->getSender() << " message: " << message << std::endl;
			if(message.find("treesize") != std::string::npos){
				int size = boost::lexical_cast<int>(message.substr(message.find(':')+1,std::string::npos));

				int side = worldModel.msgReceived[i].getSender()[0] - 'A';

				// send acknowledgement
				worldModel.msgToSend.push_back(new Message(worldModel.msgReceived[i].getSender(),"ack"));

//				std::cout << worldModel.robotId << " received size info from side: " << side << " size: " << size << std::endl;
				counted[side] = size;
			}
		}
	}else if (!done){
		for(size_t i=0;i<worldModel.msgReceived.size();i++){
			std::string message = worldModel.msgReceived[i].getDataString();

			if(message.find("organsimsize") != std::string::npos){
				int size = boost::lexical_cast<int>(message.substr(message.find(':')+1,std::string::npos));

//				std::cout << worldModel.robotId << " Organism size calculated: " << size << std::endl;
				worldModel.organismSize = size;

				boost::property_tree::ptree& supervisorTree = worldModel.additionalProperties.get_child("Supervisor");
				startTime = getTime();
				supervisorTree.put("state", 3);

				done = true;
			}
		}
	}
}

void WebotsSupervisor::setOrganismChannel(){
	// Set the channel for the organism emitter/receiver through local messages
	boost::property_tree::ptree& supervisorTree = worldModel.additionalProperties.get_child("Supervisor");
	if(worldModel.isEgg){
//		std::cout << worldModel.robotId << " organism channel: " << worldModel.organismChannel << " time: " << getTime()-startTime << " state: " << supervisorTree.get<int>("state") << std::endl;
		if(worldModel.organismChannel == -1){
			worldModel.organismChannel = rand();

			robot->getOrganismReceiver()->setChannel(worldModel.organismChannel);
			robot->getOrganismEmitter()->setChannel(worldModel.organismChannel);

		}else if(getTime()-startTime > ORGANISM_CHANNEL_NEGOTIATION_TIME) {
            std::cout << "Organism channel set to: " << worldModel.organismChannel << std::endl;
			startTime = getTime();
			supervisorTree.put("state", 2);
		}else{
			std::string msg = "channel:"+ boost::lexical_cast<std::string>(worldModel.organismChannel);
			worldModel.msgToSend.push_back(new Message("A",msg));
			worldModel.msgToSend.push_back(new Message("B",msg));
			worldModel.msgToSend.push_back(new Message("C",msg));
			worldModel.msgToSend.push_back(new Message("D",msg));
		}
	}else{
//        std::cout << worldModel.robotId << " organism channel: " << worldModel.organismChannel << " time: " << getTime()-startTime << " state: " << supervisorTree.get<int>("state") << std::endl;
		if(worldModel.organismChannel == -1){
			for(size_t i=0;i<worldModel.msgReceived.size();i++){
				std::string message = worldModel.msgReceived[i].getDataString();

				if(message.find("channel") != std::string::npos){
					worldModel.organismChannel = boost::lexical_cast<int>(message.substr(message.find(':')+1,std::string::npos));

					robot->getOrganismReceiver()->setChannel(worldModel.organismChannel);
					robot->getOrganismEmitter()->setChannel(worldModel.organismChannel);

					worldModel.msgToSend.push_back(new Message("A",message));
					worldModel.msgToSend.push_back(new Message("B",message));
					worldModel.msgToSend.push_back(new Message("C",message));
					worldModel.msgToSend.push_back(new Message("D",message));
				}
			}
		}else if(getTime()-startTime > ORGANISM_CHANNEL_NEGOTIATION_TIME){
			startTime = getTime();
			supervisorTree.put("state", 2);
		}
	}
}

void WebotsSupervisor::loadWorldModel() {
//	std::cout << "Updating world model" << std::endl;
	worldModel.time = getTime();

	for (size_t i = 0; i < robot->getNumDistanceSensors(); i++) {
		worldModel.irSensor[i] = floor(robot->getDistanceSensorValue(i));
	}

	for (size_t i=0;i<4;i++){
		worldModel.presence[i] = detectPresence(i);
	}
    
	worldModel.currentHingePosition = robot->getMainHinge()->getPosition();

	const double* position = robot->getGps()->getValues();
	worldModel.xPos = position[0];
	worldModel.yPos = position[2]; // Y and Z are the other way around in webots
	worldModel.zPos = position[1];


//	std::cout << "Loading connector status" << std::endl;
//	for (unsigned int i = 0; i < connector.size(); i++) {
//		worldModel.connected[i] = isConnected(i);
//	}

//	for (unsigned int msgId = 0; msgId < worldModel.msgReceived.size(); msgId++) {
//		delete worldModel.msgReceived[msgId];
//	}
//	for (unsigned int msgId = 0; msgId < worldModel.irMsgReceived.size(); msgId++) {
//		delete worldModel.irMsgReceived[msgId];
//	}
    
    receiveMessages();
}


void WebotsSupervisor::applyWorldModel() {
//	std::cout << worldModel.robotId << " setting hinge to position: " << worldModel.desiredHingePosition << std::endl; //<< " current position: " << worldModel.currentHingePosition << std::endl;
	robot->getMainHinge()->setPosition(worldModel.desiredHingePosition);
    
    robot->setLeftSpeed(worldModel.desiredSpeed[0]);
    robot->setRightSpeed(worldModel.desiredSpeed[1]);

	for (size_t i = 0; i < robot->getNumConnectors(); i++) {
		if(worldModel.toConnect[i] && !worldModel.connected[i]){
			connect(i);
		}
		if(worldModel.toDisconnect[i] && worldModel.connected[i]) {
			disconnect(i);
		}
	}
    
    for(std::map<int,int>::iterator it = worldModel.ledColours.begin(); it != worldModel.ledColours.end(); it++){
        setLEDColour(*it);
    }
    worldModel.ledColours.clear();

	sendMessages();
}


void WebotsSupervisor::receiveMessages() {
	worldModel.radioMsgReceived.clear();
	worldModel.msgReceived.clear();
	worldModel.irMsgReceived.clear();

	if(robot->getOrganismReceiver()->getChannel() != worldModel.organismChannel){
        robot->getOrganismReceiver()->setChannel(worldModel.organismChannel);
	}
	const char sides[17] = "ABCDABCDABCDABCD";

	// Receive Radio Messages
//    std::cout << "Receiving radio messages" << std::endl;
	while (robot->getRadioReceiver()->getQueueLength()) {
        std::string data((char*)robot->getRadioReceiver()->getData(),robot->getRadioReceiver()->getDataSize());
		worldModel.radioMsgReceived.push_back(new Message("*", data));
		robot->getRadioReceiver()->nextPacket();
	}

	// Receive Organism Messages
//    std::cout << "Receiving organism messages" << std::endl;
	while (robot->getOrganismReceiver()->getQueueLength()) {
        std::string data((char*)robot->getOrganismReceiver()->getData(),robot->getOrganismReceiver()->getDataSize());
		worldModel.msgReceived.push_back(new Message("", data));
		robot->getOrganismReceiver()->nextPacket();
	}

	// Receive IR Messages
//    std::cout << "Receiving IR messages" << std::endl;
	for (size_t i = 0; i < robot->getNumIRReceivers(); i++) {
		if(robot->hasIRReceiver(i)) {
			while (robot->getIRReceiver(i)->getQueueLength()) {
                if(worldModel.connected[i%8/2]){
                    std::stringstream sender;
                    sender << sides[(i%8)/2];
                    std::string data((char*)robot->getIRReceiver(i)->getData(),robot->getIRReceiver(i)->getDataSize());
                    worldModel.msgReceived.push_back(new Message(sender.str(),data));
                }else{
                    std::string data((char*)robot->getIRReceiver(i)->getData(),robot->getIRReceiver(i)->getDataSize());
                    worldModel.irMsgReceived.push_back(new IRMessage((unsigned int)i, data,robot->getIRReceiver(i)->getSignalStrength()));
                }
				robot->getIRReceiver(i)->nextPacket();
			}
		}
	}
//    std::cout << worldModel.robotId << " received " << worldModel.msgReceived.size() << " normal messages" << std::endl;
//    std::cout << worldModel.robotId << " received " << worldModel.irMsgReceived.size() << " IR messages" << std::endl;
}


void WebotsSupervisor::sendMessages() {
	if(robot->getOrganismReceiver()->getChannel() != worldModel.organismChannel){
		robot->getOrganismReceiver()->setChannel(worldModel.organismChannel);
	}

//	std::cout << worldModel.robotId << " need to send " << worldModel.msgToSend.size() << " normal messages" << std::endl;
//	std::cout << worldModel.robotId << " need to send " << worldModel.radioMsgToSend.size() << " radio messages" << std::endl;
//	std::cout << worldModel.robotId << " need to send " << worldModel.irMsgToSend.size() << " ir messages" << std::endl;

	// Emitters go through other robots, so set the emitter range
	// according to whether a side is connected or not.
	for(int i=0;i<4;i++){
		if(worldModel.connected[i]){
			if(robot->getIREmitter(2*i))
				robot->getIREmitter(2*i)->setRange(0.05);
		}else{
			if(robot->getIREmitter(2*i))
				robot->getIREmitter(2*i)->setRange(0.5);
		}
	}


	// send organism messages
	for (unsigned int i = 0; i < worldModel.msgToSend.size(); i++) {
		std::stringstream s;
		switch (worldModel.msgToSend[i].getSender().at(0)) {
			case 'A':
				if(worldModel.connected[0])
					assert(robot->getIREmitter(0)->send(worldModel.msgToSend[i].getData(), (int)worldModel.msgToSend[i].getDataLength()));
//					emitter[8]->send((void*)s.str().c_str(), s.str().length());
				break;
			case 'B':
				if(worldModel.connected[1])
					assert(robot->getIREmitter(2)->send(worldModel.msgToSend[i].getData(), (int)worldModel.msgToSend[i].getDataLength()));
//					emitter[10]->send((void*)s.str().c_str(), s.str().length());
				break;
			case 'C':
				if(worldModel.connected[2])
					assert(robot->getIREmitter(4)->send(worldModel.msgToSend[i].getData(), (int)worldModel.msgToSend[i].getDataLength()));
//					emitter[12]->send((void*)s.str().c_str(), s.str().length());
				break;
			case 'D':
				if(worldModel.connected[3])
					assert(robot->getIREmitter(6)->send(worldModel.msgToSend[i].getData(), (int)worldModel.msgToSend[i].getDataLength()));
//					emitter[14]->send((void*)s.str().c_str(), s.str().length());
				break;
			case '+':
				robot->getOrganismEmitter()->send(worldModel.msgToSend[i].getData(), (int)worldModel.msgToSend[i].getDataLength());
				break;
		}
//		delete worldModel.msgToSend[i];
	}
	worldModel.msgToSend.clear();

	// send Radio messages
	for (size_t i = 0; i < worldModel.radioMsgToSend.size(); i++) {
		robot->getRadioEmitter()->send(worldModel.radioMsgToSend[i].getData(),(unsigned int)worldModel.radioMsgToSend[i].getDataLength());
//		delete worldModel.radioMsgToSend[i];
	}
	worldModel.radioMsgToSend.clear();

//	std::cout << "Robot " << worldModel.robotId << " has " << worldModel.irMsgToSend.size() << " ir messages to send" << std::endl;
	// send IR messages
	for (unsigned int msgId = 0; msgId < worldModel.irMsgToSend.size(); msgId++) {
//		std::cout << worldModel.robotId << " going to send message to: " << worldModel.irMsgToSend[msgId]->getSender() << " connected: " << worldModel.connected[worldModel.irMsgToSend[msgId]->getSender()] << "...";
		if(robot->hasIREmitter(worldModel.irMsgToSend[msgId].getSender())){
			robot->getIREmitter(worldModel.irMsgToSend[msgId].getSender())->send(worldModel.irMsgToSend[msgId].getData(), (unsigned int)worldModel.irMsgToSend[msgId].getDataLength());
//			std::cout << " success!" << std::endl;
		}else{
//			std::cout << " failed..." << std::endl;
		}
//		delete worldModel.irMsgToSend[msgId];
	}
	worldModel.irMsgToSend.clear();
}


void WebotsSupervisor::connect(size_t id) {
//	std::cout << worldModel.robotId << " presence: " << connector[id]->getPresence() << std::endl;
	if(robot->hasConnector(id) && robot->getConnector(id)->getPresence() > 0){
		robot->getConnector(id)->lock();
		worldModel.connected[id] = true;
	}
	worldModel.toConnect[id] = false;
}


void WebotsSupervisor::disconnect(size_t id) {
	robot->getConnector(id)->unlock();
	worldModel.connected[id] = false;
	worldModel.toDisconnect[id] = false;
}


bool WebotsSupervisor::isConnected(size_t id) {
	return worldModel.connected[id];
}


bool WebotsSupervisor::detectPresence(size_t id) {
	if(robot->hasConnector(id)){
//		std::cout << getNameOutput() << " Detecting presence on side: " << id << ": " << connector[id]->getPresence() << std::endl;
		return robot->getConnector(id)->getPresence() > 0;
	}else{
		return false;
	}
}


int WebotsSupervisor::numConnected() {
	int count = 0;
	for(int i=0;i<4;i++){
		if(worldModel.connected[i]){
			count++;
		}
	}
	return count;
}



int WebotsSupervisor::numNotReceived(int counted[], int num) {
	int count = 0;
	for(int i=0;i<num;i++){
		if(counted[i] < 0){
			count++;
		}
	}
	return count;
}


std::string WebotsSupervisor::getOpenConnection() {
	std::string address;
	if(counted[0] < 0)
		address += "A";
	else if(counted[1] < 0)
		address += "B";
	else if(counted[2] < 0)
		address += "C";
	else if(counted[3] < 0)
		address += "D";

	return address;
}

int WebotsSupervisor::sum(int counted[],int num){
	int sum = 0;
	for(int i=0;i<num;i++){
		if(counted[i] >= 0){
			sum += counted[i];
		}
	}
	return sum;
}
